import os
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Dict, Any, Callable

# 可选依赖，如果未安装则相应功能不可用
try:
    import pdfplumber
    HAS_PDFPLUMBER = True
except ImportError:
    HAS_PDFPLUMBER = False

try:
    import docx
    HAS_DOCX = True
except ImportError:
    HAS_DOCX = False

try:
    import doc2text
    HAS_DOC2TEXT = True
except ImportError:
    HAS_DOC2TEXT = False

try:
    import oss2  # 引入 OSS 库
    HAS_OSS2 = True
except ImportError:
    HAS_OSS2 = False

class FileHandler:
    """
    一个用于处理文件提取和 OSS 存储的类。

    该类旨在:
    1. 从多种文件类型（PDF, DOCX, DOC, TXT）中提取纯文本内容。
    2. 将本地文件上传到 OSS（支持自动断点续传）。
    3. 将提取的文本字符串直接上传到 OSS。
    4. 从 OSS 下载文件到本地（支持自动断点续传）。
    """

    def __init__(self, oss_config: Optional[Dict[str, Any]] = None, executor: Optional[ThreadPoolExecutor] = None):
        """
        初始化文件处理器。

        :param oss_config: (可选) 用于OSS连接的配置字典。
                           必需键: 'endpoint', 'access_key', 'secret_key', 'bucket_name'
        :param executor: (可选) 线程池执行器，用于异步操作。如果为None，将创建默认执行器。
        """
        self.oss_config = oss_config
        self.bucket: Optional[oss2.Bucket] = None  # 用于存放 OSS Bucket 对象
        self._executor = executor or ThreadPoolExecutor(max_workers=4, thread_name_prefix="oss_upload")

        if self.oss_config:
            print("FileHandler 正在初始化 OSS 连接...")
            self._init_oss()
        else:
            print("FileHandler 已初始化（未配置 OSS）。")

    def _init_oss(self):
        """
        私有方法：根据配置初始化 OSS Bucket 客户端。
        """
        if not HAS_OSS2:
            raise RuntimeError("oss2 库未安装，请运行: pip install oss2")
        
        if not all(k in self.oss_config for k in ['endpoint', 'access_key', 'secret_key', 'bucket_name']):
            raise ValueError("OSS 配置不完整。必须包含 'endpoint', 'access_key', 'secret_key' 和 'bucket_name'。")
        
        try:
            auth = oss2.Auth(self.oss_config['access_key'], self.oss_config['secret_key'])
            bucket_name = self.oss_config['bucket_name']
            endpoint = self.oss_config['endpoint']
            
            # 创建 Bucket 实例
            self.bucket = oss2.Bucket(auth, endpoint, bucket_name)
            
            # (可选) 测试连接：尝试获取存储桶信息
            self.bucket.get_bucket_info()
            print(f"OSS 连接成功。Bucket: {bucket_name}")
            
        except oss2.exceptions.OssError as e:
            print(f"OSS 鉴权或连接失败: {e}")
            raise RuntimeError(f"无法连接到 OSS Bucket '{self.oss_config['bucket_name']}'. 错误: {e}")
        except Exception as e:
            print(f"OSS 初始化时发生未知错误: {e}")
            raise e

    def _check_oss_client(self):
        """
        私有方法：在使用OSS功能前检查客户端是否已初始化。
        """
        if self.bucket is None:
            raise RuntimeError("OSS 未配置或初始化失败。请在实例化 FileHandler 时提供有效的 oss_config。")

    # --- 文本提取功能 (与上一版相同) ---

    def extract_text(self, file_path: str, encoding: str = 'utf-8') -> str:
        """
        从给定路径的文件中提取文本内容。

        这是一个主分发方法，它会根据文件扩展名调用相应的私有提取方法。

        :param file_path: 文件的完整路径。
        :param encoding: 仅用于 .txt 文件的编码，默认为 'utf-8'。
        :return: 提取的文本字符串。
        :raises FileNotFoundError: 如果文件不存在。
        :raises ValueError: 如果文件类型不受支持。
        :raises RuntimeError: 如果在提取过程中发生内部错误。
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件未找到: {file_path}")

        _, file_extension = os.path.splitext(file_path)
        file_extension = file_extension.lower()

        try:
            if file_extension == '.pdf':
                if not HAS_PDFPLUMBER:
                    raise RuntimeError("pdfplumber 库未安装，无法处理 PDF 文件。请运行: pip install pdfplumber")
                return self._extract_pdf(file_path)
            elif file_extension == '.docx':
                if not HAS_DOCX:
                    raise RuntimeError("python-docx 库未安装，无法处理 DOCX 文件。请运行: pip install python-docx")
                return self._extract_docx(file_path)
            elif file_extension == '.doc':
                if not HAS_DOC2TEXT:
                    raise RuntimeError("doc2text 库未安装，无法处理 DOC 文件。请运行: pip install doc2text")
                return self._extract_doc(file_path)
            elif file_extension == '.txt':
                return self._extract_txt(file_path, encoding)
            else:
                raise ValueError(f"不支持的文件类型: {file_extension}")
                
        except Exception as e:
            print(f"处理文件 {file_path} 时出错: {e}")
            raise RuntimeError(f"提取失败: {file_path}。错误: {e}") from e

    def _extract_pdf(self, file_path: str) -> str:
        text_parts = []
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text_parts.append(page_text)
        return "\n".join(text_parts)

    def _extract_docx(self, file_path: str) -> str:
        doc = docx.Document(file_path)
        text_parts = []
        for para in doc.paragraphs:
            text_parts.append(para.text)
        return "\n".join(text_parts)

    def _extract_doc(self, file_path: str) -> str:
        doc = doc2text.doc2text(file_path)
        return doc.text

    def _extract_txt(self, file_path: str, encoding: str) -> str:
        try:
            with open(file_path, 'r', encoding=encoding) as f:
                return f.read()
        except UnicodeDecodeError:
            try:
                with open(file_path, 'r', encoding='gbk') as f:
                    return f.read()
            except Exception as e:
                raise UnicodeError(f"无法使用 {encoding} 或 gbk 解码文件: {file_path}。错误: {e}")


    # --- 新增的 OSS 功能 ---

    def upload_to_oss(self, local_file_path: str, object_name: str, use_resumable: bool = True, multipart_threshold_mb: int = 100):
        """
        将本地文件上传到 OSS（同步方法）。

        :param local_file_path: 本地文件的完整路径。
        :param object_name: 存储在 OSS 上的对象名称（即路径）。
        :param use_resumable: 是否启用断点续传（分块上传）。
        :param multipart_threshold_mb: (仅当 use_resumable=True 时有效)
                                        文件大小超过多少 MB 时触发断点续传。
                                        OSS 默认的简单上传限制是 5GB，但断点续传对大文件更稳健。
        """
        self._check_oss_client()  # 检查 OSS 是否已初始化
        
        if not os.path.exists(local_file_path):
            raise FileNotFoundError(f"本地文件未找到: {local_file_path}")

        file_size = os.path.getsize(local_file_path)
        threshold_bytes = multipart_threshold_mb * 1024 * 1024

        try:
            if use_resumable and file_size > threshold_bytes:
                print(f"文件大小 ({file_size} B) 超过阈值 ({threshold_bytes} B)，使用断点续传...")
                # oss2.resumable_upload 会自动处理分片、重试和并发
                # store 参数用于记录上传断点，以便下次续传
                import tempfile
                temp_dir = tempfile.gettempdir()
                oss2.resumable_upload(
                    self.bucket,
                    object_name,
                    local_file_path,
                    store=oss2.ResumableStore(root=temp_dir), # 使用系统临时目录
                    multiget_threshold=threshold_bytes, # 内部下载分片大小
                    part_size=1024 * 1024 * 10, # 10MB 分片 (可调)
                    num_threads=4 # 并发线程数
                )
                print("断点续传完成。")
            else:
                print("使用简单上传...")
                # 对于小文件，简单上传更高效
                self.bucket.put_object_from_file(object_name, local_file_path)
                print("简单上传完成。")
                
            print(f"文件 '{local_file_path}' 已成功上传到 OSS: '{object_name}'")

        except Exception as e:
            if HAS_OSS2 and hasattr(oss2, 'exceptions') and isinstance(e, oss2.exceptions.OssError):
                print(f"OSS 上传失败 (Object: {object_name}): {e}")
                raise RuntimeError(f"OSS 上传失败: {e}")
            else:
                print(f"上传过程中发生未知错误: {e}")
                raise e

    async def upload_to_oss_async(
        self,
        local_file_path: str,
        object_name: str,
        use_resumable: bool = True,
        multipart_threshold_mb: int = 100
    ) -> str:
        """
        异步将本地文件上传到 OSS（不会阻塞主线程）。

        :param local_file_path: 本地文件的完整路径。
        :param object_name: 存储在 OSS 上的对象名称（即路径）。
        :param use_resumable: 是否启用断点续传（分块上传）。
        :param multipart_threshold_mb: 触发断点续传的文件大小阈值（MB）。
        :return: OSS 对象的访问URL。
        """
        self._check_oss_client()
        
        if not os.path.exists(local_file_path):
            raise FileNotFoundError(f"本地文件未找到: {local_file_path}")

        # 在线程池中执行同步上传操作，不阻塞事件循环
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            self._executor,
            self.upload_to_oss,
            local_file_path,
            object_name,
            use_resumable,
            multipart_threshold_mb
        )
        
        # 生成 OSS URL（根据实际需求调整URL格式）
        # 这里假设使用公共读或签名URL，具体取决于OSS配置
        try:
            # 如果bucket是公共读，可以直接构造URL
            endpoint = self.oss_config['endpoint']
            bucket_name = self.oss_config['bucket_name']
            # 移除 http:// 或 https:// 前缀（如果有）
            endpoint_clean = endpoint.replace('http://', '').replace('https://', '')
            oss_url = f"https://{bucket_name}.{endpoint_clean}/{object_name}"
            return oss_url
        except Exception:
            # 如果无法生成URL，返回对象名称
            return object_name

    def upload_text_to_oss(self, text_content: str, object_name: str):
        """
        将提取的文本字符串直接作为文件内容上传到 OSS。

        :param text_content: 要上传的字符串内容。
        :param object_name: 存储在 OSS 上的对象名称（例如 "extracted/my_doc.txt")。
        """
        self._check_oss_client()
        
        try:
            # put_object 接受 bytes 或 str。如果是 str，默认使用 utf-8 编码。
            self.bucket.put_object(object_name, text_content)
            print(f"文本内容已成功上传到 OSS: '{object_name}'")
            
        except oss2.exceptions.OssError as e:
            print(f"OSS 文本上传失败 (Object: {object_name}): {e}")
            raise RuntimeError(f"OSS 文本上传失败: {e}")
        except Exception as e:
            print(f"文本上传过程中发生未知错误: {e}")
            raise e

    def download_from_oss(self, object_name: str, local_file_path: str, use_resumable: bool = True, multipart_threshold_mb: int = 100):
        """
        从 OSS 下载文件到本地。

        :param object_name: OSS 上的对象名称。
        :param local_file_path: 要保存到的本地路径。
        :param use_resumable: 是否启用断点续传（分块下载）。
        :param multipart_threshold_mb: 触发断点续传的文件大小阈值。
        """
        self._check_oss_client()

        try:
            # 1. 检查对象是否存在
            if not self.bucket.object_exists(object_name):
                raise FileNotFoundError(f"OSS 上不存在该对象: {object_name}")

            # 2. 决定是否使用断点续传
            use_resumable_download = False
            if use_resumable:
                try:
                    # 获取文件元数据（主要是大小）
                    meta = self.bucket.head_object(object_name)
                    file_size = int(meta.headers.get('Content-Length', 0))
                    threshold_bytes = multipart_threshold_mb * 1024 * 1024
                    
                    if file_size > threshold_bytes:
                        use_resumable_download = True
                        print(f"文件大小 ({file_size} B) 超过阈值，使用断点续传...")
                except oss2.exceptions.OssError as e:
                    print(f"警告: 无法获取对象大小 ({object_name})，将使用简单下载。错误: {e}")
            
            # 确保本地目录存在
            local_dir = os.path.dirname(local_file_path)
            if local_dir and not os.path.exists(local_dir):
                os.makedirs(local_dir)

            # 3. 执行下载
            if use_resumable_download:
                oss2.resumable_download(
                    self.bucket,
                    object_name,
                    local_file_path,
                    store=oss2.ResumableStore(root='/tmp'), # 断点记录目录
                    multiget_threshold=threshold_bytes,
                    part_size=1024 * 1024 * 10, # 10MB 分片
                    num_threads=4
                )
                print("断点续传下载完成。")
            else:
                print("使用简单下载...")
                self.bucket.get_object_to_file(object_name, local_file_path)
                print("简单下载完成。")

            print(f"文件 '{object_name}' 已成功下载到: '{local_file_path}'")

        except oss2.exceptions.OssError as e:
            print(f"OSS 下载失败 (Object: {object_name}): {e}")
            raise RuntimeError(f"OSS 下载失败: {e}")
        except Exception as e:
            print(f"下载过程中发生未知错误: {e}")
            raise e